package org.robotfish.learn;

import org.ejml.simple.SimpleMatrix;

/**
 * PolyFeatures gives polynomial features, e.g., (x,y) -> (x, x^2, y, y^2) 
 * @author Jaak Simm
 */
public class PolyFeatures extends Features {
	int[] inputs;
	int[] powers;
	int constant;
	boolean sign;

	/**
	 * @param numInputs  number of inputs to use, indeces of inputs (0, ..., numInputs-1)
	 * @param powers     powers to use
	 * @param constant   whether to use a constant feature (1.0)
	 * @param sing       whether to use sign(x)*abs(x^power) or ordinary x^power
	 * PolyFeature that uses the first numInputs as features.
	 */
	public PolyFeatures(int numInputs, int[] powers, boolean constant, boolean sign) {
		this(null, powers, constant, sign);
		inputs = new int[numInputs];
		for (int i=0; i<numInputs; i++) {
			inputs[i] = i;
		}
	}

	/**
	 * @param inputs    indeces of inputs to use
	 * @param powers    powers to use
	 * @param constant  whether to use a constant feature (1.0)
	 * @param sign      whether to use sign(x)*abs(x^power) or ordinary x^power
	 */
	public PolyFeatures(int[] inputs, int[] powers, boolean constant, boolean sign) {
		this.inputs = inputs;
		this.powers = powers;
		this.constant = constant ?1 :0;
		this.sign = sign;
	}

	@Override
	public int dimensions() {
		return constant+inputs.length*powers.length;
	}

	@Override
	public double[] feat(SimpleMatrix x) {
		double[] f = new double[dimensions()];
		if (constant==1) {
			f[0] = 1.0;
		}
		if (!sign) {
			// x^power
			for (int i=0; i<inputs.length; i++) {
				for (int j=0; j<powers.length; j++) {
					f[constant+i*powers.length+j] = Math.pow( x.get(inputs[i]), powers[j] );
				}
			}
		} else {
			// sign(x) * abs(x)^power
			for (int i=0; i<inputs.length; i++) {
				for (int j=0; j<powers.length; j++) {
					f[constant+i*powers.length+j] = x.get(inputs[i]) * Math.pow( Math.abs(x.get(inputs[i])), powers[j]-1 );
				}
			}
		}
		return f;
	}

}
